// Copyright (C) Kumo inc. and its affiliates.
// Author: Jeff.li lijippy@163.com
// All rights reserved.
// This program is free software: you can redistribute it and/or modify
// it under the terms of the GNU Affero General Public License as published
// by the Free Software Foundation, either version 3 of the License, or
// (at your option) any later version.
//
// This program is distributed in the hope that it will be useful,
// but WITHOUT ANY WARRANTY; without even the implied warranty of
// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
// GNU Affero General Public License for more details.
//
// You should have received a copy of the GNU Affero General Public License
// along with this program.  If not, see <https://www.gnu.org/licenses/>.
//


#include <cstdint>
#include <optional>
#include <pollux/testing/gtest_utils.h>
#include <pollux/testing/functions/prestosql/function_base_test.h>

using namespace kumo::pollux;
using namespace kumo::pollux::test;
using namespace kumo::pollux::exec;
using namespace kumo::pollux::functions::test;

namespace {
class ArrayTrimTest : public FunctionBaseTest {};

TEST_F(ArrayTrimTest, bigintArrays) {
  auto test = [this](
                  const std::vector<int64_t>& inputArray,
                  int size,
                  const std::vector<int64_t>& expectedOutput) {
    auto input = make_array_vector<int64_t>({inputArray});
    auto result = evaluate(
        fmt::format("trim_array(c0, {})", size), make_row_vector({input}));
    auto expected = make_array_vector<int64_t>({expectedOutput});
    assertEqualVectors(expected, result);
  };

  test({1, 2, 3, 4}, 0, {1, 2, 3, 4});
  test({1, 2, 3, 4}, 1, {1, 2, 3});
  test({1, 2, 3, 4}, 2, {1, 2});
  test({1, 2, 3, 4}, 3, {1});
  test({1, 2, 3, 4}, 4, {});

  auto input = make_nullable_array_vector<int64_t>({{1, 2, 3, 4}});
  POLLUX_ASSERT_THROW(
      evaluate("trim_array(c0, 5)", make_row_vector({input})),
      "size must not exceed array cardinality. arraySize: 4, size: 5");

  POLLUX_ASSERT_THROW(
      evaluate("trim_array(c0, -1)", make_row_vector({input})),
      "size must not be negative: -1");
}

TEST_F(ArrayTrimTest, simpleIntVector) {
  auto test = [this](
                  const VectorPtr& inputArrayVector,
                  int size,
                  const VectorPtr& expectedArrayVector) {
    auto result = evaluate(
        fmt::format("trim_array(c0, {})", size),
        make_row_vector({inputArrayVector}));
    assertEqualVectors(expectedArrayVector, result);
  };

  const auto input = make_array_vector<int64_t>({{1, 2, 3, 4}, {3, 4, 5}, {6}});
  const auto expected = make_array_vector<int64_t>({{1, 2, 3}, {3, 4}, {}});
  test({input}, 1, {expected});
}

TEST_F(ArrayTrimTest, complexIntVector) {
  auto test = [this](
                  const VectorPtr& inputArrayVector,
                  int size,
                  const VectorPtr& expectedArrayVector) {
    auto result = evaluate(
        fmt::format("trim_array(c0, {})", size),
        make_row_vector({inputArrayVector}));
    assertEqualVectors(expectedArrayVector, result);
  };

  auto seedVector =
      make_array_vector<int64_t>({{1, 1}, {2, 2}, {3, 3}, {4, 4}, {5, 5}});

  // Create arrays of array vector using above seed vector.
  // [[1, 1], [2, 2], [3, 3]]
  // [[4, 4], [5, 5]]
  const auto arrayOfArrayInput = make_array_vector({0, 3}, seedVector);

  // [[1, 1], [2, 2]]
  // [[4, 4]]
  const auto expected = make_array_vector(
      {0, 2}, make_array_vector<int64_t>({{1, 1}, {2, 2}, {4, 4}}));

  test({arrayOfArrayInput}, 1, {expected});
}

TEST_F(ArrayTrimTest, varcharArraysWithNull) {
  auto input = make_nullable_array_vector<std::string>(
      {{"aa", "bb", "dd"}, {"ad", std::nullopt, std::nullopt, "de"}});
  auto result = evaluate("trim_array(c0, 2)", make_row_vector({input}));
  auto expected =
      make_nullable_array_vector<std::string>({{"aa"}, {"ad", std::nullopt}});
  assertEqualVectors(expected, result);
}

} // namespace
